#include	<x86intrin.h>
#include	<sys/time.h>
#include	<algorithm>
#include	<iostream>
#include	<unistd.h>
#include	"utf16.h"
#include	<cstring>
#include	<sstream>
#include	<cfloat>
#include	<cstdio>
#include	<vector>
#include	<bitset>
#include	<cmath>
#include	<ctime>
#include	<omp.h>
using	namespace	std;
#define	_F2C	16.0f
#define	_C2F	0.0625f
float	resnet_scaling=1;
uint64_t	prng=time(NULL);
static inline uint64_t wyrand(uint64_t	*seed) {
    *seed+=0xa0761d6478bd642full;
    uint64_t  see1=*seed^0xe7037ed1a0b428dbull;
    see1*=(see1>>32)|(see1<<32);
    return	(*seed*((*seed>>32)|(*seed<<32)))^((see1>>32)|(see1<<32));
}
static	inline float wy2gau(uint64_t r) {
    const float _wynorm=1.0f/(1ull<<20);
    return ((r&0x1fffff)+((r>>21)&0x1fffff)+((r>>42)&0x1fffff))*_wynorm-3.0f;
}
static inline double wy2u01(uint64_t r) {
    const double _wynorm=1.0/(1ull<<52);
    return (r>>12)*_wynorm;
}
inline void wymum(uint64_t *A,	uint64_t *B) {
    uint64_t	hh=(*A>>32)*(*B>>32), hl=(*A>>32)*(uint32_t)*B, lh=(uint32_t)*A*(*B>>32), ll=(uint64_t)(uint32_t)*A*(uint32_t)*B;
    *A=((hl>>32)|(hl<<32))^hh;
    *B=((lh>>32)|(lh<<32))^ll;
}
inline	uint64_t	wyhash64(uint64_t	A,	uint64_t	B) {
    A^=0xa0761d6478bd642full;
    B^=0xe7037ed1a0b428dbull;
    wymum(&A,&B);
    A^=0xe7037ed1a0b428dbull;
    B^=0xa0761d6478bd642full;
    wymum(&A,&B);
    return	A^B;
}

template<class	T,	uint32_t	N>
struct	tensor {
    T	*p;
    tensor() {
        p=(T*)aligned_alloc(64,N*sizeof(T));
        memset(p,0,N*sizeof(T));
    }
//	tensor(){	p=(T*)_aligned_malloc(N*sizeof(T),64);	}
    ~tensor() {
        free(p);
    }
//	~tensor(){	_aligned_free(p);	}
    bool	load(FILE	*F) {
        return	fread(p,N*sizeof(T),1,F)==1;
    }

    void	zero(void) {
        memset(p,0,N*sizeof(T));
    }
};

template<uint32_t	R,	uint32_t	C>
struct	simd8 {
    float	scale[C];
    tensor<int8_t,R*C>	m;
    void	float2simd(float	*p,	uint32_t	c) {
        float	ma=0;
        for(uint32_t	i=0;	i<R;	i++)	ma=max(ma,fabsf(p[i]));
        scale[c]=ma=127/ma;
        int8_t	*q=m.p+c*R;
        for(uint32_t	i=0;	i<R;	i++)	q[i]=roundf(ma*p[i]);
    }
};

template<uint32_t	R,	uint32_t	C,	uint32_t	S>
struct	simd8t {
    const	float	scale=127/sqrtf(S);
    tensor<int8_t,R*C>	m;
    void	float2simd(float	*p,	uint32_t	c) {
        for(uint32_t	i=0;	i<R;	i++)	m.p[i*C+c]=roundf(scale*p[i]);
    }
};

template<uint32_t	R,	uint32_t	C>
struct	simd16 {
    float	scale[C];
    tensor<int16_t,R*C>	m;
    void	float2simd(float	*p,	uint32_t	c) {
        float	ma=0;
        for(uint32_t	i=0;	i<R;	i++)	ma=max(ma,fabsf(p[i]));
        scale[c]=ma=min(((1u<<31)-1)/ma/127/R,32767/ma);
        int16_t	*q=m.p+c*R;
        for(uint32_t	i=0;	i<R;	i++)	q[i]=roundf(ma*p[i]);
    }
};

template<uint32_t	R>
int32_t	dot(int16_t	*p,	int8_t	*q) {
    int32_t	s=0;
    for(uint32_t	i=0;	i<R;	i++)	s+=(int32_t)p[i]*(int32_t)q[i];
    return	s;
}

template<uint32_t	R0,	uint32_t	R1>
struct	linear {
    static	simd16<R0,1>	inq;
    simd8<R0,R1>	wei;
    tensor<float,R1>	out;
    void	load(FILE	*F) {
        wei.m.load(F);
    }
    void	fw(tensor<float,R0>	&inp) {
        inq.float2simd(inp.p,0);
        float	alf=1/sqrtf(R0)/inq.scale[0]/_F2C;
        #pragma omp parallel for
        for(uint32_t	i=0;	i<R1;	i++)	out.p[i]=alf*dot<R0>(inq.m.p,wei.m.p+i*R0);
    }
};
template<uint32_t	R0,	uint32_t	R1>
simd16<R0,1>	linear<R0,R1>::inq;

template<uint32_t	R,	uint32_t	H>
void	rmsnorm(tensor<float,R>	&inp) {
    uint32_t	r=R/H;
    for(uint32_t	h=0;	h<H;	h++) {
        float	nor=0,	*p=inp.p+h*r;
        for(uint32_t	i=0;	i<r;	i++)	nor+=p[i]*p[i];
        nor=sqrtf(r/nor);
        for(uint32_t	i=0;	i<r;	i++)	p[i]*=nor;
    }
}
template<uint32_t	R>
void	softmax(tensor<float,R>	&inp) {
    double	ma=-FLT_MAX,	sum=0;
    for(uint32_t	i=0;	i<R;	i++)	if(inp.p[i]>ma)	ma=inp.p[i];
    for(uint32_t	i=0;	i<R;	i++)	sum+=(inp.p[i]=expf(inp.p[i]-ma));
    for(uint32_t	i=0;	i<R;	i++)	inp.p[i]/=sum;
}

template<uint32_t	R,	uint32_t	C,	uint32_t	H,	uint32_t	S>
struct	gemat {
    tensor<float,4*H*S>	r,s;
    linear<R,4*H*S>	x;
    linear<H*S,R>	o;
    simd16<S,H>	vq;
    simd8<S,C>	vk[H];
    simd8t<S*H,C,S>	vv;
    simd16<C,H>	qa;
    tensor<float,H*S>	va;
    void	load(FILE	*F) {
        x.load(F);
        tensor<int8_t,4*H*S>	t;
        t.load(F);
        for(uint32_t	i=0;	i<4*H*S;	i++)	r.p[i]=1/(1+expf(_C2F*t.p[i]));
        o.load(F);
    }
    void	fw(tensor<float,R>	&inp,	uint32_t	col,	uint32_t	para) {
        x.fw(inp);
        for(uint32_t	i=0;	i<4*H*S;	i++)	x.out.p[i]=(s.p[i]+=r.p[i]*(x.out.p[i]-s.p[i]));
        rmsnorm<4*H*S,4*H>(x.out);
        vv.float2simd(x.out.p+2*H*S,col);
        #pragma omp parallel for
        for(uint32_t	h=0;	h<H;	h++) {
            vk[h].float2simd(x.out.p+h*S,col);
            vq.float2simd(x.out.p+H*S+h*S,h);
            float	alf=2/sqrtf(S),	a[C]= {},	ma=-FLT_MAX;
            for(uint32_t	j=para;	j<C;	j++) {
                uint32_t	i=(j+1+col)%C;
                ma=fmaxf(ma,a[i]=dot<S>(vq.m.p+h*S,vk[h].m.p+i*S)*alf/vq.scale[h]/vk[h].scale[i]);
            }
            for(uint32_t	j=para;	j<C;	j++) {
                uint32_t	i=(j+1+col)%C;
                a[i]=expf(a[i]-ma);
            }
            qa.float2simd(a,h);
            for(uint32_t	i=0;	i<S;	i++)	va.p[h*S+i]=dot<C>(qa.m.p+h*C,vv.m.p+(h*S+i)*C)/qa.scale[h]/vv.scale;
        }
        rmsnorm<H*S,H>(va);
        for(uint32_t	i=0;	i<H*S;	i++)	va.p[i]*=x.out.p[3*H*S+i];
        o.fw(va);
        for(uint32_t	i=0;	i<R;	i++)	inp.p[i]=resnet_scaling*inp.p[i]+o.out.p[i];
    }
};

template<uint32_t	R,	uint32_t	C,	uint32_t	H,	uint32_t	S>
struct	gema {
    tensor<float,2*H*S>	r,s;
    linear<R,2*H*S>	x;
    linear<H*S,R>	o;
    tensor<float,H*S>	va;
    void	load(FILE	*F) {
        x.load(F);
        tensor<int8_t,2*H*S>	t;
        t.load(F);
        for(uint32_t	i=0;	i<2*H*S;	i++)	r.p[i]=1/(1+expf(_C2F*t.p[i]));
        o.load(F);
    }
    void	fw(tensor<float,R>	&inp) {
        x.fw(inp);
        for(uint32_t	i=0;	i<2*H*S;	i++)	x.out.p[i]=(s.p[i]+=r.p[i]*(x.out.p[i]-s.p[i]));
        rmsnorm<2*H*S,2*H>(x.out);
        for(uint32_t	i=0;	i<H*S;	i++)	va.p[i]=x.out.p[i]*x.out.p[H*S+i];
        o.fw(va);
        for(uint32_t	i=0;	i<R;	i++)	inp.p[i]=resnet_scaling*inp.p[i]+o.out.p[i];
    }
};

template<uint32_t	R,	uint32_t	C,	uint32_t	H,	uint32_t	S>
struct	wyGPT {
    gema<R,C,H,S>	a;
    gema<R,C,H,S>	b;
    gemat<R,C,H,S>	c;
    gema<R,C,H,S>	d;
    gema<R,C,H,S>	e;
    void	load(FILE	*F) {
        a.load(F);
        b.load(F);
        c.load(F);
        d.load(F);
        e.load(F);
    }
    void	fw(tensor<float,R>	&inp,	uint32_t	col,	uint32_t	para) {
        a.fw(inp);
        b.fw(inp);
        c.fw(inp,col,para);
        d.fw(inp);
        e.fw(inp);
    }
};
struct	Sort {
    uint32_t	i;
    float	p;
    bool	operator()(Sort	X,	Sort	Y) {
        return	X.p>Y.p;
    }
};
typedef	bitset<65536>	mask_t;
template<uint32_t	C,	uint32_t	E,	uint32_t	D,	uint32_t	H,	uint32_t	S,	uint32_t	O>
struct	HomoSapiens {
    uint32_t	curr=0,	para=C-1;
    uint16_t	code[O],table[65536];
    vector<uint16_t>	his;
    tensor<float,E>	emb;
    vector<wyGPT<E,C,H,S>	>	tra;
    linear<E,O>	out;
    Sort	vs[O];
    HomoSapiens() {
        resnet_scaling=expf(-1.0f/(5*D));
    }
    void	string2mask(string	s,	mask_t	&m) {
        vector<uint16_t>	v;
        v.resize(s.size());
        v.resize(utf8_to_utf16((uint8_t*)s.data(),s.size(),v.data(),v.size()));
        for(size_t	i=0;	i<v.size();	i++)	m.set(v[i]);
    }
    void	range2mask(uint32_t	s,	uint32_t	e,	mask_t	&m) {
        for(uint32_t	i=s;	i<e;	i++)	m.set(i);
    }
    bool	load(const	char	*F) {
        tra.resize(D);
        FILE* f=fopen(F, "rb");
        if(f==NULL)	return	false;
        uint32_t	x;
        if(fread(&x,4,1,f)!=1||x!=C)	fprintf(stderr,"C=%u\n",x);
        if(fread(&x,4,1,f)!=1||x!=E)	fprintf(stderr,"E=%u\n",x);
        if(fread(&x,4,1,f)!=1||x!=D)	fprintf(stderr,"D=%u\n",x);
        if(fread(&x,4,1,f)!=1||x!=H)	fprintf(stderr,"H=%u\n",x);
        if(fread(&x,4,1,f)!=1||x!=S)	fprintf(stderr,"S=%u\n",x);
        if(fread(&x,4,1,f)!=1||x!=O)	fprintf(stderr,"O=%u\n",x);
        if(fread(code,O*2,1,f)!=1)	return	false;
        for(uint32_t i=0; i<D; i++)	tra[i].load(f);
        out.load(f);
        fclose(f);
        for(size_t	i=0;	i<65536;	i++)	table[i]=0;
        for(size_t	i=0;	i<O;	i++)	table[code[i]]=i;
        return	true;
    }
    void	reset(void) {
        his.clear();
        curr=0;
        para=C-1;
        for(uint32_t	i=0;	i<D;	i++) {
            tra[i].a.s.zero();
            tra[i].b.s.zero();
        }
    }
    uint16_t	sample(mask_t	&forbid) {
        uint32_t	para=his.size()<=C?C-his.size():0;
        for(uint32_t	r=0;	r<E;	r++)	emb.p[r]=(wyhash64(*his.rbegin(),r)&1)*2-1.0f;
        for(uint32_t	d=0;	d<D;	d++)	tra[d].fw(emb,curr,para);
        curr=(curr+1)%C;
        rmsnorm<E,1>(emb);
        out.fw(emb);
        for(uint32_t    i=0;    i<O;    i++)	out.out.p[i]=forbid[code[i]]?-1000:out.out.p[i]-vs[i].p;
        softmax<O>(out.out);
        double	ent=0;
        for(uint32_t    i=0;    i<O;    i++) {
            vs[i].i=i;
            vs[i].p=out.out.p[i];
            if(vs[i].p>0)	ent-=vs[i].p*logf(vs[i].p);
        }
        sort(vs,vs+O,Sort());
        uint32_t	n=0;
        double	sum=0,	plnp=0,	diversity=1/M_SQRT2;
        do {
            sum+=vs[n].p;
            plnp+=vs[n].p>0?vs[n].p*logf(vs[n].p):0;
            n++;
        } while(n<O&&logf(sum)-plnp/sum<diversity*ent);
        double  ran=wy2u01(wyrand(&prng))*sum,  sum1=0;
        uint16_t        ret=0;
        for(size_t      i=0;    i<O;    i++) {
            sum1+=vs[i].p;
            if(sum1>=ran) {
                ret=vs[i].i;
                break;
            }
        }
        return	code[ret];
    }
    void	reply(string	inp,	mask_t	&forbid,	mask_t	&stop,	uint32_t	nstop) {
        vector<uint16_t>	v;
        v.resize(inp.size());
        v.resize(utf8_to_utf16((uint8_t*)inp.data(),inp.size(),v.data(),v.size()));
        uint16_t	c;
        char	s[4];
        for(uint32_t	i=0;	i<v.size();	i++) {
            uint32_t	m=utf16_to_utf8(&v[i],1,(uint8_t*)s,4);
            for(uint32_t	j=0;	j<m;	j++)	cout<<s[j];
            cout.flush();
            his.push_back(v[i]);
            memset(vs,0,sizeof(Sort)*O);
            for(size_t	k=0;	k<his.size();	k++) {
                uint32_t	l=1;
                while(l<=k&&his[k-l]==his[his.size()-l])	l++;
                vs[table[his[k]]].p+=(float)l/(his.size()-k);
            }
            c=sample(forbid);
        }
        while(his.size()<nstop) {
            uint32_t	m=utf16_to_utf8(&c,1,(uint8_t*)s,4);
            for(uint32_t	j=0;	j<m;	j++)	cout<<s[j];
            cout.flush();
            his.push_back(c);
            memset(vs,0,sizeof(Sort)*O);
            for(size_t	k=0;	k<his.size();	k++) {
                uint32_t	l=1;
                while(l<=k&&his[k-l]==his[his.size()-l])	l++;
                vs[table[his[k]]].p+=(float)l/(his.size()-k);
            }
            c=sample(forbid);
        }
    }
};
#ifndef	WEIXIN
#include	"config"
int	main(int	ac,	char	**av) {
    HomoSapiens<context,embed,depth,headn,heads,voca>	model;
    string	model_file="model";
    int	opt,sens=context,threads=max(omp_get_num_procs()/2,1);
    while((opt=getopt(ac,	av,	"m:n:t:"))>=0) {
        switch(opt) {
        case	'm':
            model_file=optarg;
            break;
        case	'n':
            sens=atoi(optarg);
            break;
        case	't':
            threads=atoi(optarg);
            break;
        }
    }
    omp_set_num_threads(threads);
    if(!model.load(model_file.c_str())) {
        fprintf(stderr,"fail to load %s\n",model_file.c_str());
        return	0;
    }
    mask_t	forbid,stop;
    model.range2mask(0,256,forbid);
    model.string2mask("《》“”",forbid);
    model.string2mask("。！？；",stop);
    timeval	beg,end;
    gettimeofday(&beg,NULL);
    model.reset();
    model.reply(av[optind],forbid,stop,sens);
    cout<<'\n';
    gettimeofday(&end,NULL);
    cerr<<end.tv_sec-beg.tv_sec+1e-6*(end.tv_usec-beg.tv_usec)<<'\n';
    return	0;
}
#else
#include	"mongoose.h"
#include	"config"
using	namespace	std;
static struct mg_serve_http_opts s_http_server_opts;
HomoSapiens<context,embed,depth,headn,heads,voca>	model;
unordered_map<string,string>	ma;
static void ev_handler(struct mg_connection *nc, int ev, void *p) {
    if (ev != MG_EV_HTTP_REQUEST)	return;
    struct http_message	*hm=(struct http_message *) p;
    const	char	*beg=strstr(hm->body.p,"<Content><![CDATA[");
    const	char	*end=strstr(hm->body.p,"]]></Content>");
    if(beg==NULL||end==NULL)	return;
    string	query;
    for(const	char	*p=beg+strlen("<Content><![CDATA[");	p<end;	p++)	query.push_back(*p);
    beg=strstr(hm->body.p,"<FromUserName><![CDATA[");
    end=strstr(hm->body.p,"]]></FromUserName>");
    if(beg==NULL||end==NULL)	return;
    string	user;
    for(const   char    *p=beg+strlen("<FromUserName><![CDATA[");    p<end;  p++)    user.push_back(*p);
    beg=strstr(hm->body.p,"<ToUserName><![CDATA[");
    end=strstr(hm->body.p,"]]></ToUserName>");
    if(beg==NULL||end==NULL)	return;
    string	myid;
    for(const   char    *p=beg+strlen("<ToUserName><![CDATA[");    p<end;  p++)    myid.push_back(*p);
    beg=strstr(hm->body.p,"<MsgId>");
    end=strstr(hm->body.p,"</MsgId>");
    if(beg==NULL||end==NULL)	return;
    string	mid;
    for(const   char    *p=beg+strlen("<MsgId>");    p<end;  p++)    mid.push_back(*p);
    string	ret;
    auto	mi=ma.find(mid);
    if(mi!=ma.end())	ret=mi->second;
    else {
        timeval	bt,et;
        gettimeofday(&bt,NULL);
        ret=model.generate(query,672);
        ma[mid]=ret;
        gettimeofday(&et,NULL);
        cerr<<et.tv_sec-bt.tv_sec+1e-6*(et.tv_usec-bt.tv_usec)<<'\n';
    }
    ostringstream	so;
    so.precision(4);
    so.setf(ios::fixed);
    so<<"<xml>\n";
    so<<"<ToUserName><![CDATA["<<user<<"]]></ToUserName>\n";
    so<<"<FromUserName><![CDATA["<<myid<<"]]></FromUserName>\n";
    so<<"<CreateTime>"<<time(NULL)<<"</CreateTime>\n";
    so<<"<MsgType><![CDATA[text]]></MsgType>\n";
    so<<"<Content><![CDATA["<<ret<<"]]></Content>\n";
    so<<"</xml>";
    mg_send_head(nc, 200, so.str().size(), "");
    mg_printf(nc,	"%s",	so.str().c_str());
    mg_serve_http(nc, (struct http_message *) p, s_http_server_opts);
}
int main(int	ac,	char	**av) {
    if(ac<2) {
        cerr<<"weixin port\n";
        return	0;
    }
    model.load("model");
    struct mg_mgr mgr;
    struct mg_connection *nc;
    mg_mgr_init(&mgr, NULL);
    nc = mg_bind(&mgr, av[1], ev_handler);
    mg_set_protocol_http_websocket(nc);
    for(;;)	mg_mgr_poll(&mgr, 15000);
    mg_mgr_free(&mgr);
    return 0;
}
#endif

